#pragma once

#include "impl/sddmm_kernels.cuh"
#include "impl/gat_kernels.cuh"

template <typename Idx, typename DType>
void FusedGATImpl( CSRInfo<Idx> csr, GATData<Idx, DType> gdata ) {
  const Idx feat_dim = gdata.feat_src_hidden;
  const Idx h_dim = gdata.feat_src_xlen;
  const Idx e_heads = gdata.heads;
  const int cta_size = 256;

  int ntx, nty, nbx, nby;
  
  constexpr int tile = 32;
  constexpr bool fuse = false;
  using coop = llvlib::CoopLite<tile>;

  ntx = coop::thread_x();
  nty = coop::thread_y(cta_size);
  const dim3 blk(ntx, nty);
  nbx = coop::block_x(csr.rows, cta_size);

  SWITCH_4(e_heads, head, {
    SWITCH_4(MIN(4, h_dim/tile), unroll, {
      nby = coop::block_y(h_dim, 1, unroll);
      const dim3 grd(nbx, nby);
      CUDA_LAUNCH_CHECK(
        gatAllFusedv2<Idx, DType, head, tile, unroll><<<grd, blk>>>(gdata, csr);
      );
    });  
  });
}


template <typename Idx, typename DType, int HEADS, int TILE, int UNROLL>
void gat_exp_leakyrelu_sum(Idx *row_index, 
                     Idx *row_offset, 
                     Idx *col_indx, 
                     DType *edge_weight, 
                     DType *src,
                     const uint ld_src,
                     DType *dst_exp,
                     DType *dst_sum,
                     const Idx nv, const Idx ne, const Idx vlen)
{
  const Idx e_heads = ld_src;
  const int cta_size = THD;

  int ntx, nty, nbx, nby;
  
  using coop = llvlib::CoopLite<TILE>;
  const int by = THD/TILE;
  const dim3 blk(TILE, by);
  const dim3 grd(CEIL(nv, by));

  GATData<Idx, DType> gdata{
    0, //feat_src_xlen
    0, //feat_src_hidden
    ld_src, // heads
    0, //ret_xlen
    0.157, //leaky_rely_slope
    col_indx, // eids
    nullptr, src, src, // feat_src, el, er
    dst_sum, dst_exp // sum, exp
  };

  CSRInfo<Idx> csr{
    nv, //rows
    ne, //nnz
    row_offset,
    col_indx,
    row_index
  };

  CUDA_LAUNCH_CHECK(
    gatExpLeakyReluSumKernel_vect<<<grd, blk>>>(gdata, csr);
  );

}